﻿using System;
using System.Collections.Generic;
using System.Linq;

public static class IEnumerableExtensions
{
    /// <summary>
    /// 从集合中随机选择一个元素（高效版本，不强制转换整个集合）
    /// </summary>
    /// <typeparam name="T">元素类型</typeparam>
    /// <param name="source">源集合</param>
    /// <returns>随机选择的元素，如果集合为空则返回默认值</returns>
    public static T RandomElement<T>(this IEnumerable<T> source)
    {
        if (source == null)
            throw new ArgumentNullException(nameof(source));

        // 优化：对于ICollection<T>可以直接获取Count
        if (source is ICollection<T> collection)
        {
            if (collection.Count == 0)
                return default(T);

            int index = UnityEngine.Random.Range(0, collection.Count);

            // 进一步优化：对于IList<T>可以直接索引访问
            if (collection is IList<T> list)
                return list[index];

            // 对于其他ICollection<T>，使用ElementAt
            return source.ElementAt(index);
        }

        // 对于未知的IEnumerable，使用 reservoir sampling 算法
        return RandomElementIterator(source);
    }

    /// <summary>
    /// 使用蓄水池抽样算法从流式数据中随机选择元素
    /// </summary>
    private static T RandomElementIterator<T>(IEnumerable<T> source)
    {
        T result = default(T);
        int count = 0;
        var random = new System.Random();

        foreach (var element in source)
        {
            count++;
            // 以 1/count 的概率替换当前结果
            if (random.Next(count) == 0)
            {
                result = element;
            }
        }

        if (count == 0)
            return default(T);

        return result;
    }

    /// <summary>
    /// 从集合中随机选择一个元素（使用Unity随机数）
    /// </summary>
    public static T RandomElementUnity<T>(this IEnumerable<T> source)
    {
        if (source == null)
            throw new ArgumentNullException(nameof(source));

        if (source is ICollection<T> collection)
        {
            if (collection.Count == 0)
                return default(T);

            int index = UnityEngine.Random.Range(0, collection.Count);

            if (collection is IList<T> list)
                return list[index];

            return source.ElementAt(index);
        }

        // 使用蓄水池抽样
        T result = default(T);
        int count = 0;

        foreach (var element in source)
        {
            count++;
            // 以 1/count 的概率替换当前结果
            if (UnityEngine.Random.Range(0, count) == 0)
            {
                result = element;
            }
        }

        return count == 0 ? default(T) : result;
    }

    /// <summary>
    /// 从集合中随机选择指定数量的元素（允许重复）
    /// </summary>
    public static IEnumerable<T> RandomElements<T>(this IEnumerable<T> source, int count, bool allowDuplicates = true)
    {
        if (source == null)
            throw new ArgumentNullException(nameof(source));
        if (count < 0)
            throw new ArgumentOutOfRangeException(nameof(count));

        // 优化：对于可索引集合直接使用随机索引
        if (source is IList<T> list && list.Count > 0)
        {
            for (int i = 0; i < count; i++)
            {
                int randomIndex = UnityEngine.Random.Range(0, list.Count);
                yield return list[randomIndex];
            }
            yield break;
        }

        // 对于其他类型，缓存到数组（只在需要时）
        var array = source as T[] ?? source.ToArray();
        if (array.Length == 0)
            yield break;

        for (int i = 0; i < count; i++)
        {
            int randomIndex = UnityEngine.Random.Range(0, array.Length);
            yield return array[randomIndex];
        }
    }
}